package basic.courseLearn11;

import org.junit.Test;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Stack;

//并查集
public class UnionFindSet<K> {

    //元素映射
    private HashMap<K, Node> elementMap;
    //集合大小
    private HashMap<Node, Integer> sizeMap;
    //父节点映射
    private HashMap<Node,Node> fatherMap;

    //节点类
    public class Node<K> {
        K value;

        public Node() {
        }

        public Node(K value) {
            this.value = value;
        }

        public K getValue() {
            return value;
        }

        public void setValue(K value) {
            this.value = value;
        }
    }

    //添加集合
    public void addElement(K k){
        if(elementMap.containsKey(k)){
            return;
        }
        Node node = new Node(k);
        elementMap.put(k, node);
        sizeMap.put(node, 1);
        fatherMap.put(node, node);
    }

    //初始化并查集
    public void initSet(List<K> list){
        elementMap = new HashMap<>();
        sizeMap = new HashMap<>();
        fatherMap = new HashMap<>();
        for(K k : list){
            Node node = new Node(k);
            elementMap.put(k, node);
            sizeMap.put(node, 1);
            fatherMap.put(node, node);
        }

    }

    //判断元素是否在并查集中
    public boolean isSameSet(K a, K b){
        //1.判断a，b是否是并查集中的元素
        if(!elementMap.containsKey(a) ||!elementMap.containsKey(b)){
            return false;
        }
        //2.找到a，b的根节点，同时进行扁平化处理
        Node rootA = findFather(a);
        Node rootB = findFather(b);
        //3.判断根节点是否相同
        return rootA == rootB;
    }

    //合并两个集合
    public void unionSet(K a, K b){
        //1. 判断a，b是否是并查集中的元素，判断a，b是否属于同一个集合
        boolean isSame = isSameSet(a, b);
        boolean test1 = elementMap.containsKey(a);
        boolean test2 = elementMap.containsKey(b);
        if(!elementMap.containsKey(a) || !elementMap.containsKey(b) || isSameSet(a,b)){
            return;
        }
        //2. 找到a，b的根节点，同时进行扁平化处理
        Node rootA = findFather(a);
        Node rootB = findFather(b);
        //3. 比较根节点的集合大小，小的集合合并到大的集合中
        int sizeA = sizeMap.get(rootA);
        int sizeB = sizeMap.get(rootB);
        Node bigRoot = sizeA > sizeB? rootA : rootB;
        Node smallRoot = bigRoot == rootA? rootB : rootA;
        //4. 更新节点映射，小的集合的父节点映射到大的集合的根节点
        fatherMap.put(smallRoot, bigRoot);
        //5. 更新集合大小
        sizeMap.put(bigRoot, sizeA + sizeB);
        //6. 删除小的集合
        sizeMap.remove(smallRoot);
    }

    //找到元素a的根节点，同时进行扁平化处理
    private Node findFather(K a) {
        //1. 找到元素a的节点
        Node node = elementMap.get(a);
        //2. 定义栈
        Stack<Node> stack = new Stack<>();
        //3. 进行循环，将路径上的所有节点压入栈中
        while(true){
            //3.1 将node压入栈中
            stack.push(node);
            //3.2 找到node的父节点
            node = fatherMap.get(node);
            //3.3 如果node是和栈中的节点相同，则跳出循环
            if(node == stack.peek()){
                break;
            }
        }
        //4. 进行扁平化处理，将路径上的所有节点的父节点都指向根节点
        Node fatherNode = stack.pop();
        while(!stack.isEmpty()){
            //4.1 获得栈顶节点
            Node topNode = stack.pop();
            //4.2 更新栈顶节点的父节点
            fatherMap.put(topNode, fatherNode);
        }
        return fatherNode;
    }

    @Test
    public void test(){
        List<Integer> nodes=new ArrayList<>();
        for (int i = 0; i < 10; i++) {
            nodes.add(i);
        }
        UnionFindSet<Integer> unionFindSet = new UnionFindSet<>();
        unionFindSet.initSet(nodes);
        unionFindSet.unionSet(nodes.get(0), nodes.get(1));
        unionFindSet.unionSet(nodes.get(2), nodes.get(3));
        unionFindSet.unionSet(nodes.get(4), nodes.get(5));
        unionFindSet.unionSet(nodes.get(6), nodes.get(7));
        unionFindSet.unionSet(nodes.get(8), nodes.get(9));
        System.out.println(unionFindSet.isSameSet(nodes.get(0), nodes.get(1))); //true
        System.out.println(unionFindSet.isSameSet(nodes.get(2), nodes.get(3))); //true
        System.out.println(unionFindSet.isSameSet(nodes.get(4), nodes.get(5))); //true
        System.out.println(unionFindSet.isSameSet(nodes.get(6), nodes.get(7))); //true
        System.out.println(unionFindSet.isSameSet(nodes.get(8), nodes.get(9))); //true
        System.out.println(unionFindSet.isSameSet(nodes.get(0), nodes.get(2))); //false
        System.out.println(unionFindSet.isSameSet(nodes.get(1), nodes.get(3))); //false
        System.out.println(unionFindSet.isSameSet(nodes.get(4), nodes.get(6))); //false
        System.out.println(unionFindSet.isSameSet(nodes.get(5), nodes.get(7))); //false
        System.out.println(unionFindSet.isSameSet(nodes.get(8), nodes.get(0))); //false
    }

}
